import os
from CalTorque import CalTorque
from CalTqFixOffset_Gauss import CorrectForOffset_G
from CalTqFixOffset_PeakFinder import CorrectForOffset_P
from datetime import datetime
import matplotlib
from matplotlib import pyplot as PLT
from scipy.optimize import curve_fit
from numpy import std

data_files = []
dp = 'data/'
ls= os.listdir(dp)
for i in ls:
   if i.split('.')[-1] not in ['xlsx', 'xls']: continue
   if 'out' in i: continue
   elif 'in' not in i: continue
   data_files.append(dp+i)

p0 = []
offset_G = []

filelist = []

for f in data_files:   
   data = CalTorque(f)
   if data['source'] != 'weakTh-228': continue
   if data['station'] != 'S5': continue
   if data['position'][0] > 50: continue
   if max(data['pos_offsetfix']) < 350: continue
   print f
   
   x0 = data['position'][0]
   p0.append(x0)
   filelist.append(f.split('.')[0].split(' ')[2].split('_')[0])
      
   offset_G.append(data['offset'])


x = list(xrange(-25, 15))
y = x
y1 = []
y2 = []
for item in y:
    y1.append(item + 1)
    y2.append(item - 1)
    

fig1 = PLT.figure(figsize=(15,8),dpi=150)
box = [0.14, 0.14, 0.76, 0.76]
ax1 = fig1.add_axes(box)
ax1.set_ylabel('Initial Position Offset (cm)')
ax1.set_xlabel('Poly-Two-Gauss Offset (cm)')
PLT.title('Initial Position Versus Poly-Two-Gauss Offset)')
ax1.annotate('y = x', xy = (15, 10), color = 'r', fontstyle = 'italic', fontsize = 'x-large')
ax1.plot(x, y1, '-', color = 'g')
ax1.plot(x, y2, '-', color = 'g')
ax1.scatter(offset_G, p0, color = 'maroon')
#for i, txt in enumerate(filelist):
    #ax1.annotate(txt, (offset_G[i], p0[i]), textcoords = 'offset points', ha = 'right', va = 'bottom', size = 'x-small', alpha = 0.5, rotation = 45)
ax1.plot(x, y, '-', color = 'r', linewidth = 2)
fig1.savefig('p0_vs_Offset_Gauss.png')

thislist = []

for i, thing in enumerate(p0):
    this = thing - offset_G[i]
    thislist.append(this)

fig2 = PLT.figure(figsize = (15,8), dpi = 150)
ax2 = fig2.add_axes(box)
ax2.set_ylabel('Initial Position spread')
ax2.set_xlabel('Data Offset (cm)')
PLT.title('Initial Position Spread versus Offset (model method)')
ax2.scatter(offset_G, thislist, color = 'r')
fig2.savefig('p0_vs_Offset_Gauss_spread.png')

stan = std(thislist)

# stan = 1.27457602256

